# %%

import os
import json
import numpy as np
import matplotlib.pyplot as plt

# 初始化项
data_folder = r"C:\py\hrv_predict\data_hrv"
age_min = 0
age_max = 100
age_all_group = 10
# gender_limit = '男'
gender_limit = '无'
if gender_limit == '男':
    graph_append = 'and Gender: Male'
elif gender_limit == '女':
    graph_append = 'and Gender: Female'
else:
    graph_append = ''


# 读取全部文件存入参数字典
hrv_dict = {}
hrv_group_dict = {}
age_data = []
age_space = (age_max - age_min) // age_all_group

for filename in os.listdir(data_folder):
    if filename.endswith('.mat'):
        record = filename.split('.')[0]
        file_1 = os.path.join(data_folder, record + '.json')
        file_2 = os.path.join(data_folder, record + '_hrv.json')
        if os.path.exists(file_1) and os.path.exists(file_2):
            with open(file_2, 'r') as hrv_file:
                hrv_data = json.load(hrv_file)
            with open(file_1, 'r') as user_file:
                user_data = json.load(user_file)
                analysis = user_data['AnalysisText']
                age = user_data['age']
                gender = user_data['Gender']
            if (graph_append and gender == gender_limit) or (not graph_append):
                if analysis and hrv_data['nbeats'] > 40000 and 0 < age < 100:
                    if '颤' not in analysis and \
                            '停' not in analysis and \
                            '就医' not in analysis and \
                            '频' not in analysis:
                        age_data.append(age)
                        for field in hrv_data:
                            if field not in hrv_dict:
                                hrv_dict[field] = []
                            hrv_dict[field].append(hrv_data[field])
            else:
                pass


# %%
print(str(hrv_dict.keys()))
# %%
# 按照年龄分组保存字典
age_data = np.array(age_data)
print(age_data.shape)
# 年龄分组
age_group = (age_data // age_space).astype(int)
for field in hrv_dict:
    if field not in hrv_group_dict:
        hrv_group_dict[field] = {}
        for index, group in enumerate(age_group):
            if group not in hrv_group_dict[field]:
                hrv_group_dict[field][group] = []
            hrv_group_dict[field][group].append(hrv_dict[field][index])
    # print(hrv_group_dict)
    hrv_dict[field] = np.array(hrv_dict[field])
    for group in age_group:
        hrv_group_dict[field][group] = np.array(hrv_group_dict[field][group])

# %%
print(hrv_group_dict['mean_nni'])
# %%

# # 年龄分布直方图
# plt.hist(age_data, bins=120, edgecolor='k')
# plt.title('Age Distribution')
# plt.xlabel('Age')
# plt.ylabel('Frequency')
# plt.grid(True)
# plt.show()


# # 查看变量散点图
# for i in hrv_dict:
#     print(i)
#     x = age_data
#     y = hrv_dict[i]
#     point_size = 3
#     plt.scatter(x, y, point_size)
#     title = hrv_dict.items()
#     plt.title(i)
#     plt.show()
# %%
# 性别限制下各指标的各年龄组均值和方差
cal_dict = {field: {group: {'mean_value': [], 'variance_value': []} for group in range(age_all_group)} for field in hrv_dict}

for field in hrv_group_dict:
    if field == 'tinn':
        break
    for group in age_group:
        y = hrv_group_dict[field][group]
        if not cal_dict[field][group]['mean_value']:
            mean_value = np.mean(y)
            variance_value = np.var(y)
            cal_dict[field][group]['mean_value'] = mean_value
            cal_dict[field][group]['variance_value'] = variance_value
# for field in cal_dict:
#     print(field)
#     print(cal_dict[field])
# %%
# 性别限制下各指标分年龄组箱型图
for field in hrv_group_dict:
    if field == 'tinn':
        break
    fig, axes = plt.subplots(nrows=1, ncols=age_all_group, sharex=True, sharey=True)
    for group in age_group:
        y = hrv_group_dict[field][group]
        ax = plt.subplot(1, age_all_group, group + 1)
        label = f'{group * age_space}-{group * age_space + age_space}'
        ax.set_xlabel(label)
        ax.boxplot(y)
        ax.set_xticklabels([])

    plt.suptitle(f'Boxplot of {field} Values by Age {graph_append}')
    plt.tight_layout()
    plt.show()
 
# %%
# # 绘制箱型图同时显示对应均值和方差
# cal_dict = {field: {group: {'mean_value': [], 'variance_value': []} for group in range(age_all_group)} for field in
#             hrv_dict}

# for field in hrv_group_dict:
#     if field == 'tinn':
#         break
#     fig, axes = plt.subplots(nrows=1, ncols=age_all_group, sharex=True, sharey=True)
#     for group in age_group:
#         x = group
#         y = hrv_group_dict[field][group]
#         ax = plt.subplot(1, age_all_group, group + 1)
#         # ax.axis('off')
#         label = f'{group * age_space}-{group * age_space + age_space}'
#         ax.set_xlabel(label)
#         ax.boxplot(y)
#         ax.set_xticklabels([])

#         if cal_dict[field][group]['mean_value']:
#             pass
#         else:
#             mean_value = np.mean(y)
#             variance_value = np.var(y)
#             # print(f"Field: {field}, Group: {group}, Mean: {mean_value}, Variance: {variance_value}")
#             cal_dict[field][group]['mean_value'] = mean_value
#             cal_dict[field][group]['variance_value'] = variance_value
#     # gender_limit限制下的field指标下的年龄分组
#     # 每组均值和方差
#     print(field)
#     print(cal_dict[field])
#     # 每组箱型图
#     plt.suptitle(f'Boxplot of {field} Values by Age {graph_append}')
#     plt.tight_layout()
#     plt.show()

# %%
